﻿#include "DoseStatistics.h"

#include <algorithm>
#include <cmath>
#include <iostream>
#include <numeric>
#include <stdexcept>
#include <string>

// 计算DVH（剂量体积直方图）
DoseStatisticsResult accumulate(const std::vector<double>& voxelVolume, const std::vector<double>& voxelDose, bool doExactAccumulation)
{
    printf("doExactAccumulation: %d\n", doExactAccumulation);
    DoseStatisticsResult result;

    // 过滤掉体积为0的体素
    std::vector<double> filteredVolume;
    std::vector<double> filteredDose;
    for (size_t i = 0; i < voxelVolume.size(); ++i)
    {
        if (voxelVolume[i] > 0)
        {
            filteredVolume.push_back(voxelVolume[i]);
            filteredDose.push_back(voxelDose[i]);
        }
    }

    size_t num = filteredDose.size();
    double num2 = 0.0;

    if (num < 5000)
    {
        // 精确计算DVH
        result.binnedDoseValues.resize(num + 2, 0.0);
        result.accumulatedRelativeVolumes.resize(num + 2, 0.0);

        for (size_t i = 0; i < num; ++i)
        {
            result.accumulatedRelativeVolumes[i + 1] = filteredVolume[i];
            result.binnedDoseValues[i + 1] = std::max(0.0, filteredDose[i]);  // 确保剂量不为负
        }

        // 排序：先按剂量排序，剂量相同时按体积排序
        std::vector<size_t> indices(num);
        std::iota(indices.begin(), indices.end(), 1);  // 从1开始，跳过第一个和最后一个元素

        std::sort(indices.begin(), indices.end(),
                  [&result](size_t a, size_t b)
                  {
                      if (result.binnedDoseValues[a] != result.binnedDoseValues[b])
                      {
                          return result.binnedDoseValues[a] < result.binnedDoseValues[b];
                      }
                      return result.accumulatedRelativeVolumes[a] < result.accumulatedRelativeVolumes[b];
                  });

        // 应用排序
        std::vector<double> sortedDoseValues(num);
        std::vector<double> sortedVolumes(num);
        for (size_t i = 0; i < num; ++i)
        {
            sortedDoseValues[i] = result.binnedDoseValues[indices[i]];
            sortedVolumes[i] = result.accumulatedRelativeVolumes[indices[i]];
        }

        // 复制回结果
        for (size_t i = 0; i < num; ++i)
        {
            result.binnedDoseValues[i + 1] = sortedDoseValues[i];
            result.accumulatedRelativeVolumes[i + 1] = sortedVolumes[i];
        }
    }
    else
    {
        // 粗略计算DVH
        double minVoxelDose = *std::min_element(filteredDose.begin(), filteredDose.end());
        double maxVoxelDose = *std::max_element(filteredDose.begin(), filteredDose.end());
        minVoxelDose = std::max(0.0, minVoxelDose);

        int splitNum = static_cast<int>(std::min(1001.0, std::max(2.0, (maxVoxelDose - minVoxelDose) / 1e-6 + 1.0)));
        if (splitNum <= 1) splitNum = 2;

        result.binnedDoseValues.resize(splitNum + 2, 0.0);
        result.accumulatedRelativeVolumes.resize(splitNum + 2, 0.0);

        if (maxVoxelDose > 1e-6)
        {
            double intervalVoxelDose = (maxVoxelDose - minVoxelDose) / (splitNum - 1);

            for (int j = 0; j < splitNum; ++j)
            {
                result.binnedDoseValues[j + 1] = minVoxelDose + j * intervalVoxelDose;
            }

            for (size_t k = 0; k < num; ++k)
            {
                int intervalIndex;
                if (filteredDose[k] > minVoxelDose)
                {
                    intervalIndex = static_cast<int>((filteredDose[k] - minVoxelDose) / intervalVoxelDose + 0.5);
                }
                else
                {
                    intervalIndex = 0;
                }
                result.accumulatedRelativeVolumes[intervalIndex + 1] += filteredVolume[k];
            }
        }
    }

    // 计算累积相对体积
    for (int i = static_cast<int>(result.accumulatedRelativeVolumes.size()) - 2; i > 0; --i)
    {
        double num9 = result.accumulatedRelativeVolumes[i];
        result.accumulatedRelativeVolumes[i] = num2 + 0.5 * num9;
        num2 += num9;
    }

    result.accumulatedRelativeVolumes[0] = num2;
    result.accumulatedRelativeVolumes.back() = 0.0;
    result.binnedDoseValues[0] = result.binnedDoseValues[1];
    result.binnedDoseValues.back() = result.binnedDoseValues[result.binnedDoseValues.size() - 2];

    // 归一化
    if (num2 > 0.0)
    {
        for (auto& volume : result.accumulatedRelativeVolumes)
        {
            volume /= num2;
        }
    }

    return result;
}

// 获取特定体积百分比对应的剂量
double get_metric(double doseAtVolume, const std::vector<double>& binnedDoseValues, const std::vector<double>& accumulatedRelativeVolumes)
{
    // 查找精确匹配
    for (size_t i = 0; i < accumulatedRelativeVolumes.size(); ++i)
    {
        if (std::abs(accumulatedRelativeVolumes[i] - doseAtVolume) < 1e-12)
        {
            printf("Accurate, doseAtVolume: %f, binnedDoseValues: %f, accumulatedRelativeVolumes: %f\n", doseAtVolume, binnedDoseValues[i],
                   accumulatedRelativeVolumes[i]);
            return binnedDoseValues[i];
        }
    }

    //printf("Coarse search\n");
    // 查找第一个大于doseAtVolume的索引
    size_t ind = 0;
    for (; ind < accumulatedRelativeVolumes.size(); ++ind)
    {
        if (accumulatedRelativeVolumes[ind] < doseAtVolume)
        {
            break;
        }
    }

    //printf("coarse target dose -> %f, index -> %zu\n", doseAtVolume, ind);
    // 边界检查
    if (ind == 0) return binnedDoseValues[0];
    if (ind >= accumulatedRelativeVolumes.size()) return binnedDoseValues.back();

    size_t ind1 = ind - 1;
    size_t ind2 = ind;

    // 线性插值
    if (std::abs(accumulatedRelativeVolumes[ind2] - accumulatedRelativeVolumes[ind1]) > 1e-10)
    {
        //printf("Linear interpolation; dose -> %f, index -> %zu\n", doseAtVolume, ind);
        double metric = (binnedDoseValues[ind2] - binnedDoseValues[ind1]) / (accumulatedRelativeVolumes[ind2] - accumulatedRelativeVolumes[ind1]) *
                            (doseAtVolume - accumulatedRelativeVolumes[ind1]) +
                        binnedDoseValues[ind1];
        return metric;
    }
    else
    {
        //printf("Average; dose -> %f, index -> %zu\n", doseAtVolume, ind);
        // 如果体积相同，返回平均剂量
        return (binnedDoseValues[ind2] + binnedDoseValues[ind1]) * 0.5;
    }
}

// 计算ROI内的剂量统计指标
double get_dose_statistics_roi_interp(const std::string& dose_type,
                                      const std::vector<double>& binnedDoseValues,
                                      const std::vector<double>& accumulatedRelativeVolumes,
                                      const std::vector<double>& voxelVolume,
                                      const std::vector<double>& voxelDose,
                                      double doseAtVolume)
{
    if (dose_type == "average")
    {
        double totalVolume = 0.0;
        double weightedDose = 0.0;

        for (size_t i = 0; i < voxelVolume.size(); ++i)
        {
            totalVolume += voxelVolume[i];
            weightedDose += voxelDose[i] * voxelVolume[i];
        }

        if (totalVolume <= 0.0)
        {
            throw std::runtime_error("Prescription ROI volume is zero. Cannot scale to average dose.");
        }

        return weightedDose / totalVolume;
    }
    else
    {
        double targetDoseAtVolume;

        if (dose_type == "D_98")
        {
            targetDoseAtVolume = 0.98;
        }
        else if (dose_type == "D_95")
        {
            targetDoseAtVolume = 0.95;
        }
        else if (dose_type == "D_50")
        {
            targetDoseAtVolume = 0.5;
        }
        else if (dose_type == "D_2")
        {
            targetDoseAtVolume = 0.02;
        }
        else if (dose_type == "DoseAtVolume")
        {
            targetDoseAtVolume = doseAtVolume;
        }
        else
        {
            throw std::runtime_error("Prescription type not supported. Only 'average', 'D_98', 'D_95', 'D_50', 'D_2', 'DoseAtVolume' are supported");
        }

        return get_metric(targetDoseAtVolume, binnedDoseValues, accumulatedRelativeVolumes);
    }
}